# src/local_vector_store.py
import logging
import os
import pickle
from typing import List, Dict, Any, Tuple

from scipy.spatial.distance import cosine


class VectorClient:
    """
    使用本地 pickle 文件来存储和检索向量数据。
    """

    def __init__(self, storage_path: str):
        self.storage_path = storage_path
        self.data = self._load_data()

    def _load_data(self) -> Dict[str, Dict[int, Any]]:
        """
        从 pickle 文件加载数据。如果文件不存在，返回一个空字典。
        """
        if os.path.exists(self.storage_path):
            try:
                with open(self.storage_path, 'rb') as f:
                    logging.info(f"Loading vectors from {self.storage_path}")
                    return pickle.load(f)
            except (pickle.UnpicklingError, EOFError) as e:
                logging.error(f"Could not load vector store file (it might be corrupted or empty): {e}")
        return {}

    def _save_data(self):
        """
        将内存中的数据保存到 pickle 文件。
        """
        os.makedirs(os.path.dirname(self.storage_path), exist_ok=True)
        with open(self.storage_path, 'wb') as f:
            pickle.dump(self.data, f)
        logging.info(f"Vector store saved to {self.storage_path}")

    def upsert_vectors(self, collection_name: str, ids: List[int], vectors: List[List[float]],
                       payloads: List[Dict[str, Any]]):
        """
        批量上传或更新向量。
        """
        if collection_name not in self.data:
            self.data[collection_name] = {}

        for i, point_id in enumerate(ids):
            self.data[collection_name][point_id] = {
                "vector": vectors[i],
                "payload": payloads[i]
            }

        self._save_data()
        logging.info(f"Upserted {len(ids)} vectors into collection '{collection_name}' in local store.")

    def search(self, collection_name: str, query_vector: List[float], top_k: int = 1) -> List[Tuple[int, float]]:
        """
        Searches the specified collection for the top_k most similar points to the query vector.
        Returns a list of tuples [(point_id, similarity_score), ...]
        """
        if collection_name not in self.data or not self.data[collection_name]:
            return []

        collection = self.data[collection_name]

        results = []
        for point_id, point_data in collection.items():
            stored_vector = point_data.get("vector")
            if stored_vector:
                # This line will now work correctly
                similarity = 1 - cosine(query_vector, stored_vector)
                results.append((point_id, similarity))

        results.sort(key=lambda x: x[1], reverse=True)

        return results[:top_k]
